#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Dec 31 22:41:47 2021

@author: liqingsimac
"""

###############################################
##条件期望
import numpy as np
import matplotlib.pyplot as plt

x=np.array([0,1])
y=np.array([0,1])
yex=np.array([0.5,0.5])

fig=plt.figure()
ax=fig.add_subplot(121)

ax.hlines(y=0,xmin=0,xmax=1)
ax.vlines(x=0,ymin=0,ymax=1)

ax.plot(x,y,'--',label=r'$X$')
ax.plot(x,yex,'-',label=r'$E(X)$')

ax.set_xlabel('omega')
ax.set_ylabel(r'$X$ and $E(X)$')
ax.legend(loc='upper left')
ax.set_xticks([0,1/2,1])
ax.set_yticks([0,1/2,1])

x1=np.array([0,1/3])
y1=np.array([1/6,1/6])
x2=np.array([1/3,2/3])
y2=np.array([1/2,1/2])
x3=np.array([2/3,1])
y3=np.array([5/6,5/6])

bx=fig.add_subplot(122)

bx.hlines(y=0,xmin=0,xmax=1)
bx.vlines(x=0,ymin=0,ymax=1)

bx.plot(x,y,'--',label=r'$X$')
bx.plot(x1,y1,'-',label=r'$E(X|B_1)$')
bx.plot(x2,y2,'-',label=r'$E(X|B_2)$')
bx.plot(x3,y3,'-',label=r'$E(X|B_3)$')
bx.set_xticks([0,1/3,2/3,1])
bx.set_yticks([0,1/6,1/2,5/6,1])

bx.set_xlabel('omega')
bx.set_ylabel(r'$X$ and $E(X|B_i)$')
bx.legend(loc='upper left')

fig.tight_layout(w_pad=2)
#fig.savefig('conditional-expectation.png')






